#version 430 core

#extension GL_NV_shader_atomic_float : require

#define GRID_SIZE_X 64
#define GRID_SIZE_Y 64

const float elementMatrix[4][4] = {
	{8.0/3.0, -2.0/3.0, -4.0/3.0, -2.0/3.0},
	{-2.0/3.0, 8.0/3.0, -2.0/3.0, -4.0/3.0},
	{-4.0/3.0, -2.0/3.0, 8.0/3.0, -2.0/3.0},
	{-2.0/3.0, -4.0/3.0, -2.0/3.0, 8.0/3.0}
};

const float rhsMatrix[4][8] = {
	{ -2.0/3.0, 2.0/3.0, 1.0/3.0, -1.0/3.0, -2.0/3.0, -1.0/3.0, 1.0/3.0, 2.0/3.0 },
	{ -2.0/3.0, 2.0/3.0, 1.0/3.0, -1.0/3.0, -1.0/3.0, -2.0/3.0, 2.0/3.0, 1.0/3.0 },
	{ -1.0/3.0, 1.0/3.0, 2.0/3.0, -2.0/3.0, -1.0/3.0, -2.0/3.0, 2.0/3.0, 1.0/3.0 },
	{ -1.0/3.0, 1.0/3.0, 2.0/3.0, -2.0/3.0, -2.0/3.0, -1.0/3.0, 1.0/3.0, 2.0/3.0 }
};


layout(local_size_x=128) in;

layout(binding=1, r32i) uniform readonly iimage2D nodeIndexImage;
layout(binding=2, rg32f) uniform readonly image2D velocityImage;

layout(std430, binding=0) coherent buffer femElementPositions {
	ivec2 elementPositions[];
};

layout(std430, binding=1) coherent buffer sparseMatrixBuffer {
	float sparseMatrix[];
};

layout(std430, binding=2) coherent buffer bVectorBuffer {
	float bVector[];
};

uniform uint numFemElements;

void main(void)
{
	uint elementID = gl_GlobalInvocationID.x;
	if (elementID < numFemElements)
	{
		ivec2 elemPos = elementPositions[elementID];

		vec2 vel0 = imageLoad(velocityImage,elemPos).xy;
		vec2 vel1 = imageLoad(velocityImage,elemPos+ivec2(1,0)).xy;
		vec2 vel2 = imageLoad(velocityImage,elemPos+ivec2(1,1)).xy;
		vec2 vel3 = imageLoad(velocityImage,elemPos+ivec2(0,1)).xy;

		int nodeIndex0 = imageLoad(nodeIndexImage,elemPos).x;
		int nodeIndex1 = imageLoad(nodeIndexImage,elemPos+ivec2(1,0)).x;
		int nodeIndex2 = imageLoad(nodeIndexImage,elemPos+ivec2(1,1)).x;
		int nodeIndex3 = imageLoad(nodeIndexImage,elemPos+ivec2(0,1)).x;

		if (nodeIndex0 >= 0)
		{
			atomicAdd(sparseMatrix[nodeIndex0*9  ], elementMatrix[0][0]);
			
			float rhs0 = rhsMatrix[0][0] * vel0.x + rhsMatrix[0][1] * vel1.x + rhsMatrix[0][2] * vel2.x + rhsMatrix[0][3] * vel3.x + 
				rhsMatrix[0][4] * vel0.y + rhsMatrix[0][5] * vel1.y + rhsMatrix[0][6] * vel2.y + rhsMatrix[0][7] * vel3.y;

			atomicAdd(bVector[nodeIndex0], rhs0);

			if (nodeIndex1 >= 0)
			{
				atomicAdd(sparseMatrix[nodeIndex0*9+4], elementMatrix[0][1]);
				atomicAdd(sparseMatrix[nodeIndex1*9+8], elementMatrix[1][0]);
			}
			if (nodeIndex2 >= 0)
			{
				atomicAdd(sparseMatrix[nodeIndex0*9+5], elementMatrix[0][2]);
				atomicAdd(sparseMatrix[nodeIndex2*9+1], elementMatrix[2][0]);
			}
			if (nodeIndex3 >= 0)
			{
				atomicAdd(sparseMatrix[nodeIndex0*9+6], elementMatrix[0][3]);
				atomicAdd(sparseMatrix[nodeIndex3*9+2], elementMatrix[3][0]);
			}
		}

		if (nodeIndex1 >= 0)
		{
			atomicAdd(sparseMatrix[nodeIndex1*9  ], elementMatrix[1][1]);

			float rhs1 = rhsMatrix[1][0] * vel0.x + rhsMatrix[1][1] * vel1.x + rhsMatrix[1][2] * vel2.x + rhsMatrix[1][3] * vel3.x + 
				rhsMatrix[1][4] * vel0.y + rhsMatrix[1][5] * vel1.y + rhsMatrix[1][6] * vel2.y + rhsMatrix[1][7] * vel3.y;

			atomicAdd(bVector[nodeIndex1], rhs1);

			if (nodeIndex2 >= 0)
			{
				atomicAdd(sparseMatrix[nodeIndex1*9+6], elementMatrix[1][2]);
				atomicAdd(sparseMatrix[nodeIndex2*9+2], elementMatrix[2][1]);
			}
			if (nodeIndex3 >= 0)
			{
				atomicAdd(sparseMatrix[nodeIndex1*9+7], elementMatrix[1][3]);
				atomicAdd(sparseMatrix[nodeIndex3*9+3], elementMatrix[3][1]);
			}
		}

		if (nodeIndex2 >= 0)
		{
			atomicAdd(sparseMatrix[nodeIndex2*9  ], elementMatrix[2][2]);

			float rhs2 = rhsMatrix[2][0] * vel0.x + rhsMatrix[2][1] * vel1.x + rhsMatrix[2][2] * vel2.x + rhsMatrix[2][3] * vel3.x + 
				rhsMatrix[2][4] * vel0.y + rhsMatrix[2][5] * vel1.y + rhsMatrix[2][6] * vel2.y + rhsMatrix[2][7] * vel3.y;

			atomicAdd(bVector[nodeIndex2], rhs2);

			if (nodeIndex3 >= 0)
			{
				atomicAdd(sparseMatrix[nodeIndex2*9+8], elementMatrix[2][3]);
				atomicAdd(sparseMatrix[nodeIndex3*9+4], elementMatrix[3][2]);
			}
		}

		if (nodeIndex3 >= 0)
		{
			atomicAdd(sparseMatrix[nodeIndex3*9 ], elementMatrix[3][3]);

			float rhs3 = rhsMatrix[3][0] * vel0.x + rhsMatrix[3][1] * vel1.x + rhsMatrix[3][2] * vel2.x + rhsMatrix[3][3] * vel3.x + 
				rhsMatrix[3][4] * vel0.y + rhsMatrix[3][5] * vel1.y + rhsMatrix[3][6] * vel2.y + rhsMatrix[3][7] * vel3.y;

			atomicAdd(bVector[nodeIndex3], rhs3);
		}
		
	}
}